"""
基于 LlamaIndex 的 RAG 系统
使用 ChromaDB 向量数据库，支持 JSON 格式的 PDF 知识库
"""

import os
import json
from typing import List, Dict
from llama_index.core import (
    VectorStoreIndex,
    Document,
    StorageContext,
    Settings
)
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.schema import NodeWithScore
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
import pypdf


# 配置 LlamaIndex 全局设置
Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",
    device="cpu"
)
Settings.chunk_size = 512
Settings.chunk_overlap = 50

# PDF 路径配置
PDF_DIR = "./pdf"
CHROMA_DB_DIR = "./chroma_db"
COLLECTION_NAME = "exercise_knowledge"

# 主题到 PDF 文件的映射
TOPIC_TO_PDF = {
    "减脂": "减脂.pdf",
    "增肌": "增肌.pdf"
}


def load_json_from_pdf(pdf_path: str) -> Dict:
    """
    从 PDF 中提取 JSON 数据
    
    Args:
        pdf_path: PDF 文件路径
        
    Returns:
        解析后的 JSON 数据
    """
    reader = pypdf.PdfReader(pdf_path)
    full_text = ""
    
    for page in reader.pages:
        full_text += page.extract_text()
    
    # 清理文本：移除 PDF 提取产生的多余换行符和空格
    # 保留 JSON 结构中的换行，但移除内容中间的换行
    full_text = full_text.replace('\n', ' ')
    
    # 移除可能的 markdown 代码块标记
    full_text = full_text.strip()
    if full_text.endswith('```'):
        full_text = full_text[:-3].strip()
    if full_text.startswith('```'):
        full_text = full_text[3:].strip()
    
    # 解析 JSON
    try:
        data = json.loads(full_text)
        return data
    except json.JSONDecodeError as e:
        print(f"Error parsing JSON from {pdf_path}: {e}")
        print(f"Text preview: {full_text[:200]}")
        return {"principles": []}


def create_documents_from_json(json_data: Dict, source_file: str) -> List[Document]:
    """
    将 JSON 数据转换为 LlamaIndex Document 对象
    
    Args:
        json_data: JSON 格式的数据
        source_file: 源文件名
        
    Returns:
        Document 对象列表
    """
    documents = []
    principles = json_data.get("principles", [])
    
    for principle in principles:
        # 构建文档内容
        content = f"标题: {principle.get('title', '')}\n\n{principle.get('content', '')}"
        
        # 构建元数据
        metadata = {
            "source": source_file,
            "principle_id": principle.get("id", ""),
            "title": principle.get("title", ""),
            "tags": ", ".join(principle.get("tags", []))
        }
        
        # 创建 Document
        doc = Document(
            text=content,
            metadata=metadata,
            id_=f"{source_file}_{principle.get('id', '')}"
        )
        documents.append(doc)
    
    return documents


def build_index(force_rebuild: bool = False) -> VectorStoreIndex:
    """
    构建或加载向量索引（使用 ChromaDB）
    
    Args:
        force_rebuild: 是否强制重建索引
        
    Returns:
        VectorStoreIndex 对象
    """
    # 初始化 ChromaDB 客户端
    chroma_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
    
    # 如果强制重建，先删除旧的 collection
    if force_rebuild:
        try:
            chroma_client.delete_collection(name=COLLECTION_NAME)
            print(f"已删除旧的 collection: {COLLECTION_NAME}")
        except Exception:
            pass
    
    # 获取或创建 collection
    try:
        chroma_collection = chroma_client.get_collection(name=COLLECTION_NAME)
        
        # 如果 collection 已存在且不强制重建，直接加载
        if not force_rebuild and chroma_collection.count() > 0:
            print(f"加载已有索引从 ChromaDB (collection: {COLLECTION_NAME})...")
            vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
            index = VectorStoreIndex.from_vector_store(vector_store)
            print(f"索引加载完成！共 {chroma_collection.count()} 个文档")
            return index
    except Exception:
        pass
    
    # 创建新的 collection
    chroma_collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME)
    
    # 构建新索引
    print("构建新的向量索引（使用 ChromaDB）...")
    all_documents = []
    
    # 遍历所有 PDF 文件
    for topic, pdf_file in TOPIC_TO_PDF.items():
        pdf_path = os.path.join(PDF_DIR, pdf_file)
        
        if not os.path.exists(pdf_path):
            print(f"警告: {pdf_path} 不存在，跳过")
            continue
        
        print(f"加载 {pdf_file}...")
        json_data = load_json_from_pdf(pdf_path)
        documents = create_documents_from_json(json_data, pdf_file)
        all_documents.extend(documents)
        print(f"  加载了 {len(documents)} 个原则")
    
    print(f"总共加载了 {len(all_documents)} 个文档")
    
    # 创建 ChromaDB 向量存储
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    
    # 创建索引
    index = VectorStoreIndex.from_documents(
        all_documents,
        storage_context=storage_context,
        show_progress=True
    )
    
    print(f"索引已保存到 ChromaDB ({CHROMA_DB_DIR})")
    
    return index


class RAGRetriever:
    """RAG 检索器，支持按主题过滤的混合检索"""
    
    def __init__(self, index: VectorStoreIndex):
        """
        初始化检索器
        
        Args:
            index: 向量索引
        """
        self.index = index
        self.retriever = VectorIndexRetriever(
            index=index,
            similarity_top_k=15  # 初始检索更多候选，提高召回率
        )
    
    def retrieve(
        self,
        query: str,
        topic: str,
        k: int = 5
    ) -> List[NodeWithScore]:
        """
        执行检索
        
        Args:
            query: 查询文本
            topic: 主题（"减脂" 或 "增肌"）
            k: 返回结果数量
            
        Returns:
            检索到的节点列表
        """
        # 获取目标 PDF 文件名
        target_pdf = TOPIC_TO_PDF.get(topic)
        if not target_pdf:
            print(f"警告: 未知主题 {topic}")
            return []
        
        # 执行检索
        nodes = self.retriever.retrieve(query)
        
        # 过滤：只保留目标主题的文档
        filtered_nodes = [
            node for node in nodes
            if node.node.metadata.get("source") == target_pdf
        ]
        
        # 返回 top-k 结果
        return filtered_nodes[:k]


def retrieve(query: str, topic: str, k: int = 5) -> List:
    """
    便捷的检索函数（为了兼容 assessment.py）
    
    Args:
        query: 查询文本
        topic: 主题
        k: 返回结果数量
        
    Returns:
        检索结果列表（模拟 LangChain Document 格式）
    """
    # 使用全局检索器（避免重复加载）
    retriever = get_retriever(force_rebuild=False)
    
    # 执行检索
    nodes = retriever.retrieve(query, topic, k)
    
    # 转换为兼容格式
    class MockDocument:
        def __init__(self, node):
            self.page_content = node.node.text
            self.metadata = node.node.metadata
    
    return [MockDocument(node) for node in nodes]


# 全局检索器实例（单例模式）
_global_retriever = None


def get_retriever(force_rebuild: bool = False) -> RAGRetriever:
    """
    获取全局检索器实例
    
    Args:
        force_rebuild: 是否强制重建索引
        
    Returns:
        RAGRetriever 实例
    """
    global _global_retriever
    
    if _global_retriever is None or force_rebuild:
        index = build_index(force_rebuild=force_rebuild)
        _global_retriever = RAGRetriever(index)
    
    return _global_retriever


if __name__ == "__main__":
    print("=" * 80)
    print("初始化 LlamaIndex RAG 系统")
    print("=" * 80)
    
    # 构建索引
    index = build_index(force_rebuild=True)
    
    # 创建检索器
    retriever = RAGRetriever(index)

    # 测试检索
    print("\n" + "=" * 80)
    print("测试检索功能")
    print("=" * 80)
    
    test_queries = [
        ("减脂的基本原理是什么", "减脂"),
        ("如何制定增肌训练计划", "增肌"),
    ]
    
    for query, topic in test_queries:
        print(f"\n查询: {query}")
        print(f"主题: {topic}")
        print("-" * 80)
        
        results = retriever.retrieve(query, topic, k=3)
        
        for i, node in enumerate(results, 1):
            print(f"\n结果 {i}:")
            print(f"  标题: {node.node.metadata.get('title', 'N/A')}")
            print(f"  ID: {node.node.metadata.get('principle_id', 'N/A')}")
            print(f"  标签: {node.node.metadata.get('tags', 'N/A')}")
            print(f"  相似度分数: {node.score:.4f}")
            print(f"  内容预览: {node.node.text[:100]}...")
    
    print("\n" + "=" * 80)
    print("RAG 系统初始化完成！")
    print("=" * 80)

